from PIL.Image import *
from zlib import compress, decompress
from socket import (SO_RCVLOWAT, socket, AF_INET, SOCK_DGRAM, SOL_SOCKET, SO_BROADCAST)
from time import sleep
from PIL.ImageGrab import grab

BUFFER_SIZE = 1400

class Image:
    # def __init__(self) -> None:
    #     self.image
    #     self.image_bytes
    #     self.image_compress
    #     self.image_decompress
    #     self.image_size

    def grabImage(self):
        self.image = grab()
        
    def setImageSize(self):
        self.image_size = self.image.size
    
    def image2bytes(self):
        # 图片转换为bytes
        self.image_bytes = self.image.tobytes()

    def bytes2compress(self):
        # 把bytes流进行压缩
        self.image_compress = compress(self.image_bytes)

    def compress2decompress(self):
        # 把压缩的字节流解压
        self.image_decompress = decompress(self.image_compress)

    def bytes2image(self):
        self.image = frombytes(mode= "RGBA", size=(2880, 1800), data=self.image_decompress, decoder_name="raw")

class SendImage:

    def __init__(self, Image, ADDR) -> None:
        self.ADDR = ADDR
        self.Image = Image
        self.sock = socket(AF_INET, SOCK_DGRAM)
        self.sock.setsockopt(SOL_SOCKET, SO_BROADCAST, True)

    def sendStart(self):
        self.sock.sendto(b'start', self.ADDR)

    def sendOver(self):
        self.sock.sendto(b'_over', self.ADDR)

    def sendImage(self):
        for i in range(len(self.Image.image_compress) // BUFFER_SIZE + 1):
            start = i * BUFFER_SIZE
            end =  start + BUFFER_SIZE
            self.sock.sendto(self.Image.image_compress[start:end], self.ADDR)
            sleep(0.0002)

class RecvImage:

    def __init__(self, ADDR) -> None:
        self.Image = Image()
        self.sock = socket(AF_INET, SOCK_DGRAM)
        self.sock.bind(ADDR)

    def recvImage(self):

        recv_data = self.sock.recvfrom(BUFFER_SIZE)#1024表示本次接收的最大字节

        send_addr = recv_data[1]
        recv_msg = recv_data[0]
        if recv_msg == b'start':
            print("Receive start")

        is_over = True
        data = b''
        while is_over:
            recv_data = self.sock.recvfrom(BUFFER_SIZE)   #1024表示本次接收的最大字节
            
            if recv_data[0].startswith(b'_over'):
                
                print("Receive over")
                is_over = False

            if is_over != False:

                data += recv_data[0]

        self.Image.image_compress = data
        return data
    
    def showImage(self):
        Image.compress2decompress()
        Image.decompress2bytes()
        img = Image.image
        img.show()



